#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Description:
import io
import socket
import time
import paramiko
from ping3 import ping
from paramiko.ssh_exception import SSHException
from scp import SCPClient

import constants
from utils.logger import Logger

log = Logger()


class SSHClient:
    def __init__(self, ip, port=constants.ENV_PORT, username="root", password=constants.ENV_PASSWORD, timeout=300):
        self.ip = ip
        self.port = port
        self.username = username
        self.password = password
        self.timeout = timeout
        self.ssh_client = None
        self.scp_client = None
        self._init_tool()

    def _init_tool(self):
        try:
            ssh_client = paramiko.SSHClient()
            ssh_client.load_system_host_keys()
            ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh_client.connect(self.ip, self.port, self.username, self.password, timeout=self.timeout,
                               allow_agent=False, look_for_keys=False)
            self.ssh_client = ssh_client
            # 创建SCP客户端
            self.scp = SCPClient(ssh_client.get_transport())
        except SSHException as e:
            raise Exception(f"SSH客户端创建失败: {e}")

    def exec_cmd(self, command, exec_timeout=None):
        """
        同步执行命令
        """
        log.info(f"command is \n {command}")
        try:
            _, std_out, std_err = self.ssh_client.exec_command(command, timeout=exec_timeout)
            result = "".join(std_out.readlines())
            while not std_out.channel.exit_status_ready():
                time.sleep(0.1)
            return result
        except socket.timeout as e:
            raise Exception("IP:%s timeout CMD:%s" % (self.ip, command))

    def exec_async_cmd(self, command) -> int:
        """
        后台执行命令，返回pid，后续可以通过pid去停止后台
        """
        try:
            bg_command = f"nohup {command} > /dev/null 2>&1 & echo $!"
            stdin, stdout, stderr = self.ssh_client.exec_command(bg_command)
            pid = int(stdout.read().decode().strip())
            return pid
        except socket.timeout as e:
            raise Exception("IP:%s timeout CMD:%s" % (self.ip, command))

    def upload_file(self, local_path: str, remote_path: str, recursive: bool = False) -> None:
        """
        上传文件/目录
        :param local_path: 本地路径
        :param remote_path: 远程路径
        :param recursive: 是否递归上传目录
        """
        self.scp.put(local_path, remote_path, recursive=recursive)
        self.scp.close()

    def download_file(self, remote_path: str, local_path: str, recursive: bool = False) -> None:
        """
        下载文件/目录
        :param remote_path: 远程路径
        :param local_path: 本地路径
        :param recursive: 是否递归下载目录
        """
        self.scp.get(remote_path, local_path, recursive=recursive)
        self.scp.close()

    def upload_template_file(self, rendered_content, remote_path):
        """
        将内存中的I/O写入到远端主机，不保留在本地
        :param rendered_content: 内容
        :param remote_path: 远端路径
        :return:
        """
        file_obj = io.StringIO(rendered_content)
        self.scp.putfo(file_obj, remote_path)
        self.scp.close()


def ping_host(ip: str, retry: int) -> bool:
    """
    使用 ping 检测主机是否存活
    :param ip: 目标主机 IP 地址
    :param retry: 重试次数
    :return: 如果主机存活，则返回 True，否则返回 False
    """
    log.info(f"开始检查 {ip} 是否可以连接...")

    while retry > 0:
        try:
            response = ping(ip, timeout=5)
            if response is not None:
                log.success(f"{ip} 已经在线，可以连接!")
                return True
            else:
                log.error(f"{ip} 无法 ping 通，正在重试... 剩余重试次数: {retry - 1}")
                retry -= 1
                time.sleep(10)
        except Exception as e:
            log.error(f"ping 错误: {e}")
            retry -= 1
            time.sleep(10)

    # 如果重试次数耗尽，返回 False
    log.error(f"无法连接 {ip}，所有重试尝试失败。")
    return False
